# Copyright 2010 Thomas Conneely
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from scipy import optimize, exp
from functools import partial

def gauss_fit(x_data, y_data, coeff_guess, noise_floor = 0):
    '''Fits an arbitary number of gaussian curves to data using a
    least squares technique (using scipy.optmize.leastsq). Requires
    initial "guesses" for the gaussian coefficients (amplitude, 
    mean and standard deviation. Returns the fitted coefficients
    and the function used for fitting.
        
    Args:
        @ x_data        - x-data to be fitted
        @ y_data        - y-data to be fitted
        @ coeff_guess   - list of guesses for gaussian coefficients
        @ noise_floor   - the noise level in the data (optional,
                          defaults to zero)

    Example:
    
    coeff, fit_func = gauss_fit(x_data, y_data, coeff_guess, noise_floor = 0):
    
    Coeff_guess is of the format [amplitude, mean, standard_dev].
    Also, the number of gaussian peaks is determined by the length 
    of coeff_guess, for n peaks there must be 3n coefficients. For
    example for two peaks (A & B) coeff_guess would be:
    
    coeff_guess = [A_amp, A_mean, A_sd, B_amp, B_mean, B_sd]
    
    where amp is amplitude and sd is standard deviation.
    '''
    
    if len(coeff) % 3 != 0:
        raise ValueError('Coeff_guess needs three coefficients for each gaussian')
        
    # calculate number of peaks (n)
    n = len(coeff) / 3
    
    # A gaussian function, that selects 3 coeffecients
    # from a list given an index offset
    def gauss(coeff, x, offset=0):
        return coeff[offset]*exp(-(x-coeff[offset+1])**2/(2*coeff[offset+2]**2))
    
    # Build list of gaussian functions one for 
    # each gaussian peak in the distribution
    fit_funcs = []
    for i in xrange(0, n):
        f = partial(gauss, offset=i*3)
        fit_funcs.append(f)
        
    # Create the fitting function, basically the sum
    # of each gaussian peak defined before
    def fit_func(coeff, x):
        ans = noise_floor
        for f in fit_funcs:
            ans += f(coeff, x)
        return ans
    
    # Error function, that calculates the delta
    # of a data point from the fitting function
    err_func = lambda coeff, x, y : fit_func(coeff, x) - y
    
    # Do the fitting
    coeff, success = optimize.leastsq(err_func, coeff_guess[:], args=(x_data, y_data))
    
    # Now create our fitted function, with coefficients defined
    fitted_func = partial(fit_func, coeff)
    
    return coeff, fitted_func
    
class GaussianFit(object):
    '''Simple class for gaussian fitting, uses gauss_fit function, but
    creates list of fitted y_data
        
    Args:
        @ x_data        - x-data to be fitted
        @ y_data        - y-data to be fitted
        @ coeff_guess   - list of guesses for gaussian coefficients
        @ noise_floor   - the noise level in the data (optional,
                          defaults to zero)
        @ x_fit         - specifies x values to calculate fitted y
                          date for (optional, if not set the x values
                          in x_data will be used)
    
    Properties:
        @ x             - supplied x data
        @ y_fitted      - y data calculated from fitting function
        @ coeff         - coefficients of fitting function
        @ fitted_func   - the fitted function
        
    See gauss_fit doc string for more details on arguments/properties
    '''

    def __init__(self, x_data, y_data, coeff_guess, noise_floor=0, x_fit=None):
        self.coeff, self.fitted_func = gauss_fit(x_data, y_data, coeff_guess, noise_floor)
        
        if x_fit:
            self.x = x_fit
        else:
            self.x = x_data
            
        self.y_fitted = self.fitted_func(self.x)
    